--[[ A simple class for hierarchical nodes with a name, parent, and children. ]]
module(..., package.seeall)

import("utils.idgen", "IDGenerator")

local id_gen = IDGenerator()


Node = class("Node")

function Node.__init(self, name, parent)
    self.name = tostring(name or id_gen(self))
    self.fullname = self.name
    self.parent = nil
    self.children = {}
    if parent ~= nil then
        self:set_parent(parent)
    end
end

function Node.__tostring(self)
    return string.format("%s(%s)", self.__name, self.fullname)
end

function Node.add_child(self, node)
    node:set_parent(self)
end

function Node.remove_child(self, node)
    if node.parent ~= self then
        error(string.format("unable to find local child node '%s'", node.name), 2)
    end
    node:set_parent(nil)
end

function Node.get_child(self, name)
    local name = tostring(name)
    local node = self.children[name]
    if node == nil then
        error(string.format("unable to find node '%s'", name), 2)
    end
    return node
end

function Node.set_parent(self, node)
    -- check for invalid arguments
    if node == self.parent then 
        return
    elseif node == self then 
        error("invalid node provided - self", 2) 
    elseif node.children[self.name] ~= nil then
        error(string.format("node name '%s' already in use at parent", node.name), 2)
    end
    -- remove self from previous parent and add it to the new parent's children table
    if self.parent ~= nil then
        self.parent.children[self.name] = nil
    end
    if node ~= nil then
        node.children[self.name] = self
    end
    self.parent = node
    self:update_fullname()
end

function Node.set_name(self, name)
    local name = tostring(name)
    if self.parent ~= nil then
        if self.parent.children[name] ~= nil then
            error(string.format("node name '%s' already in use at parent", name), 2)
        end
        self.parent.children[self.name] = nil
        self.parent.children[name] = self
    end
    self.name = name
    self:update_fullname()
end

function Node.update_fullname(self)
    if self.parent == nil then
        self.fullname = self.name
    else
        self.fullname = self.parent.fullname .. "." .. self.name
    end
    for _, node in pairs(self.children) do
        node:update_fullname()
    end
end

function Node.get_root(self)
    local node = self
    while node.parent ~= nil do
        node = node.parent
    end
    return node
end

function Node.find(self, ...)
    local node = self
    for _, name in ipairs{...} do
        node = node:get_child(name)
    end
    return node
end

function Node.show_tree(self, indent)
    indent = indent or 0
    io.write(string.rep("  ", indent), tostring(self), "\n")
    for _, node in pairs(self.children) do
        node:show_tree(indent + 1)
    end
end

function Node.propagate_down(self, funcname, ...)
    --[[ Call a function on all nodes below (and including) the current node. ]]
    self[funcname](self, ...)
    for _, node in pairs(self.children) do
        node:propagate_down(funcname, ...)
    end
end

function Node.propagate_up(self, funcname, ...)
    --[[ Call a function on all nodes from the current node (inclusive) up to the root node. ]]
    self[funcname](self, ...)
    if self.parent ~= nil then
        self.parent:propagate_up(funcname, ...)
    end
end

----------------------------------------------------------------------------------------------------
-- Test code
----------------------------------------------------------------------------------------------------
function test()
    local root = Node("root")
    
    local function add_level(node, level)
        if level > 0 then
            add_level(Node("l", node), level - 1)
            add_level(Node("r", node), level - 1)
        end
    end
    
    add_level(root, 3)
    root:show_tree()
    return root
end
